#include "label_file.h"

#include <fstream>
#include "base64.h"

LabelFileError::LabelFileError(const std::string &message)
    : std::runtime_error(message) {}

LabelFile::LabelFile(const std::string &filename) {

    if (!filename.empty()) {
        load(filename);
    }
    this->filename = filename;
}

std::string LabelFile::load_image_file(const std::string &filename) {
    cv::Mat image = cv::imread(filename, cv::IMREAD_UNCHANGED);
    if (image.empty()) {
        std::cerr << "Failed opening image file: " << filename << std::endl;
        return "";
    }

    // Apply EXIF orientation if needed (additional implementation may be required)

    // Encode image to PNG or JPEG depending on extension
    std::vector<uchar> buf;
    std::string ext = filename.substr(filename.find_last_of('.'));
    if (ext == ".jpg" || ext == ".jpeg") {
        cv::imencode(".jpg", image, buf);
    } else {
        cv::imencode(".png", image, buf);
    }
    // Convert buffer to base64 string
    std::string encoded = base64_encode(buf.data(), buf.size());
    return encoded;
}

void LabelFile::load(const std::string &filename) {
    std::ifstream f(filename);
    if (!f.is_open()) {
        throw LabelFileError("Cannot open file: " + filename);
    }

    nlohmann::json data;
    try {
        f >> data;
    } catch (nlohmann::json::parse_error& e) {
        throw LabelFileError("JSON parse error in file: " + filename);
    }

    // Keys to extract from the JSON file
    std::vector<std::string> keys = {
        "version",
        "imageData",
        "imagePath",
        "shapes",
        "flags",
        "imageHeight",
        "imageWidth",
    };
    std::vector<std::string> shape_keys = {
        "label",
        "points",
        "group_id",
        "shape_type",
        "flags",
        "description",
        "mask",
    };

    // Read imageData
    std::string imageDataStr = data.value("imageData", "");
    if (!imageDataStr.empty()) {
        // Store base64-encoded image data
        this->imageData = imageDataStr;
    } else {
        // Build image path relative to the label file
        std::string imagePath = data["imagePath"];
        std::string dir = filename.substr(0, filename.find_last_of("/\\"));
        std::string fullImagePath = dir + "/" + imagePath;
        // Load image data
        this->imageData = load_image_file(fullImagePath);
    }

    // Read flags
    auto flags_json = data.value("flags", nlohmann::json::object());
    for (auto& item : flags_json.items()) {
        this->flags[item.key()] = item.value().get<bool>();
    }

    this->imagePath = data["imagePath"];

    // Check image height and width
    int imageHeight = data.value("imageHeight", 0);
    int imageWidth = data.value("imageWidth", 0);
    _check_image_height_and_width(this->imageData, imageHeight, imageWidth);

    // Read shapes
    auto shapes_json = data["shapes"];
    for (auto& s : shapes_json) {
        Shape shape;
        shape.label = s["label"];
        shape.points = s["points"].get<std::vector<std::vector<float>>>();
        shape.shape_type = s.value("shape_type", "polygon");
        shape.flags = s.value("flags", std::map<std::string, bool>());
        shape.description = s.value("description", "");
        shape.group_id = s.value("group_id", -1);

        std::string mask_b64 = s.value("mask", "");
        if (!mask_b64.empty()) {
            // Decode mask
            std::string mask_data = base64_decode(mask_b64);
            // Convert mask data to cv::Mat
            std::vector<uchar> mask_buf(mask_data.begin(), mask_data.end());
            shape.mask = cv::imdecode(mask_buf, cv::IMREAD_GRAYSCALE);
        }

        // Other data
        for (auto& item : s.items()) {
            if (std::find(shape_keys.begin(), shape_keys.end(), item.key()) == shape_keys.end()) {
                shape.other_data[item.key()] = item.value();
            }
        }

        this->shapes.push_back(shape);
    }

    // Read other data
    for (auto& item : data.items()) {
        if (std::find(keys.begin(), keys.end(), item.key()) == keys.end()) {
            this->otherData[item.key()] = item.value();
        }
    }

    this->filename = filename;
}

void LabelFile::_check_image_height_and_width(const std::string &imageData, int &imageHeight, int &imageWidth) {
    // Decode base64 image data
    std::string image_bytes = base64_decode(imageData);
    // Convert image bytes to cv::Mat
    std::vector<uchar> img_buf(image_bytes.begin(), image_bytes.end());
    cv::Mat img = cv::imdecode(img_buf, cv::IMREAD_UNCHANGED);
    if (img.empty()) {
        throw LabelFileError("Failed to decode image data");
    }
    if (imageHeight != 0 && img.rows != imageHeight) {
        std::cerr << "imageHeight does not match with imageData or imagePath, so getting imageHeight from actual image." << std::endl;
        imageHeight = img.rows;
    }
    if (imageWidth != 0 && img.cols != imageWidth) {
        std::cerr << "imageWidth does not match with imageData or imagePath, so getting imageWidth from actual image." << std::endl;
        imageWidth = img.cols;
    }
}

void LabelFile::save(const std::string &filename, const std::vector<Shape> &shapes, const std::string &imagePath, int imageHeight, int imageWidth, const std::string &imageData, const nlohmann::json &otherData, const std::map<std::string, bool> &flags) {

    std::string encodedImageData = imageData;
    if (!imageData.empty()) {
        // Check image height and width
        _check_image_height_and_width(imageData, imageHeight, imageWidth);
    }

    nlohmann::json data;
    data["version"] = "1.0"; // Set your version here
    data["flags"] = flags;
    data["shapes"] = nlohmann::json::array();
    data["imagePath"] = imagePath;
    data["imageData"] = encodedImageData; // Base64-encoded image data
    data["imageHeight"] = imageHeight;
    data["imageWidth"] = imageWidth;

    // Add shapes
    for (const auto& shape : shapes) {
        nlohmann::json shape_json;
        shape_json["label"] = shape.label;
        shape_json["points"] = shape.points;
        shape_json["group_id"] = shape.group_id;
        shape_json["shape_type"] = shape.shape_type;
        shape_json["flags"] = shape.flags;
        shape_json["description"] = shape.description;

        if (!shape.mask.empty()) {
            // Encode mask image to base64
            std::vector<uchar> buf;
            cv::imencode(".png", shape.mask, buf);
            std::string mask_b64 = base64_encode(buf.data(), buf.size());
            shape_json["mask"] = mask_b64;
        }

        // Add other data
        for (const auto& item : shape.other_data.items()) {
            shape_json[item.key()] = item.value();
        }

        data["shapes"].push_back(shape_json);
    }

    // Add other data
    for (const auto& item : otherData.items()) {
        data[item.key()] = item.value();
    }

    // Write data to file
    try {
        std::ofstream f(filename);
        if (!f.is_open()) {
            throw LabelFileError("Cannot open file for writing: " + filename);
        }
        f << data.dump(2); // Indent by 2 spaces
        this->filename = filename;
    } catch (std::exception& e) {
        throw LabelFileError(std::string("Error saving label file: ") + e.what());
    }
}

bool LabelFile::is_label_file(const std::string &filename) {
    return filename.substr(filename.find_last_of('.')) == LabelFile::suffix;
}
